import re
from ui.parameter_frame import ParameterFrame
from ui.parameterdisplay import ParameterDisplay
from ui.operation_frame import OperationFrame

import sys
from PyQt5.QtWidgets import QApplication,  QVBoxLayout, QWidget,  QLabel
from PyQt5.QtWidgets import QMessageBox, QDialog, QVBoxLayout, QTextEdit, QScrollArea, QPushButton

import argparse
import os
import sys
from pathlib import Path


import threading
import cv2
from ultralytics import YOLO
import os


# 调用display_results更新图像
from ui.result_frame import ResultFrame


from PyQt5.QtWidgets import QFrame, QLabel, QTextEdit
from PyQt5.QtWidgets import QPushButton, QApplication
from PyQt5.QtGui import QFont
from PyQt5.QtGui import QIcon
from PyQt5.QtCore import QSize, Qt  # 导入 QSize


class BackupFrame(QFrame):
    def __init__(self, parent=None, run_detection=None):
        super().__init__(parent)
        self.setStyleSheet("background-color: white; border: 2px solid white;")
        self.setFixedSize(440, 140)
        self.setGeometry(20, 440 + 20 + 320, 440, 140)
        # 你可以在这里添加该框架的其他功能代码

        self.set_title()  # 添加标题
        self.run_detection = run_detection  # 保存传入的检测函数

        self.is_detecting = True

        # 创建按钮
        self.button1 = QPushButton("运行代码", self)
        self.button2 = QPushButton("退出系统", self)
        self.button3 = QPushButton("停止摄像", self)  # 新增停止检测按钮

        # 设置按钮的位置和大小
        self.button1.setGeometry(15, 80, 125, 50)
        self.button2.setGeometry(300, 80, 125, 50)
        self.button3.setGeometry(17+140, 80, 125, 50)  # 停止按钮位置

        # 调整按钮文字大小
        font = QFont()
        font.setPointSize(14)  # 设置文字大小为14
        self.button1.setFont(font)
        self.button2.setFont(font)
        self.button3.setFont(font)

        # 设置按钮图标
        # self.set_button_icon(self.button1, "ui/image/icons_run.png")  # 替换为实际路径
        # self.set_button_icon(self.button2, "ui/image/icons_Esc.png")  # 替换为实际路径
        # self.set_button_icon(self.button3, "ui/image/icons_Esc.png")  # 替换为停止图标

        # 设置按钮的样式（包括凸起和下陷效果）
        self.set_button_style(self.button1,  background_color="#90EE90")
        self.set_button_style(self.button2,  background_color="#FFB6C1")
        self.set_button_style(self.button3,  background_color="#FFFFE0")

        # 连接按钮的点击事件
        self.button1.clicked.connect(self.run_code)
        self.button2.clicked.connect(self.exit_system)
        self.button3.clicked.connect(self.stop_detection)  # 停止按钮的点击事件

    def set_button_icon(self, button, image_path):
        # 设置按钮的图标
        icon = QIcon(image_path)
        button.setIcon(icon)
        button.setIconSize(QSize(40, 40))  # 设置图标的大小

    def set_button_style(self, button, background_color="#f0f0f0", border_color="#ccc", pressed_color="#ddd"):
        # 设置按钮的样式
        button.setStyleSheet(f"""
            QPushButton {{
                background-color: {background_color};
                border: 2px solid {border_color};
                border-radius: 10px;
                padding: 10px 20px;
                font-size: 20px;
                font-weight: bold;
            }}
            QPushButton:pressed {{
                background-color: {pressed_color};
                border: 2px solid {border_color};
                padding: 10px 18px;  /* 点击时按钮稍微变小 */
            }}
        """)

    def set_title(self):
        title = QLabel("系统运行与退出", self)  # 设置标题文字
        title.setFont(QFont("仿宋", 18))  # 设置字体为仿宋，大小为16
        title.setAlignment(Qt.AlignCenter)  # 让文字居中
        title.setGeometry(0, 15, self.width(), 50)  # 设置位置，0为X轴，10为顶部间距，宽度为窗口宽度，高度为30
        title.setStyleSheet("color: black;  font-weight: bold;")  # 设置字体颜色为黑色

    def stop_detection(self):
        print("停止检测...")
        if self.is_detecting:
            self.is_detecting = False  # 设置为 False，停止检测


    def run_code(self):
        print("运行run代码")
        # 触发YOLO检测
        if self.run_detection:
            self.run_detection()  # 调用主窗口中的run_detection方法

    def exit_system(self):
        # 退出系统
        print("正在退出系统...")
        QApplication.quit()  # 退出应用程序



class MainWindow(QWidget):
    # update_image_signal = pyqtSignal(np.ndarray)  # 定义信号，传递图像数据

    def __init__(self):
        super().__init__()

        # 设置窗口标题
        self.setWindowTitle("YOLOv8目标检测系统")

        "设置界面颜色"
        # self.setStyleSheet("background-color: #FACAC2;")  # 设置背景颜色
        self.setStyleSheet("""
            QWidget {
                background: qlineargradient(x1: 0, y1: 0, x2: 1, y2: 0, 
                                            stop: 0 #40BCEA, stop: 0.85 #40BCEA, 
                                            stop: 0.85 #FF99FF, stop: 1 #FF99FF);
            }
        """)

        # 设置窗口大小
        # self.setGeometry(100, 100, 800, 600)
        # 调用 adjustSize() 方法使窗口大小根据内容自动调整
        # 界面大小和布局设置
        "设置界面大小的参数"
        screen = QApplication.primaryScreen()  # 获取主显示器的屏幕对象。
        screen_rect = screen.availableGeometry() # 返回屏幕的可用区域的矩形坐标，包括屏幕的宽度和高度。
        width, height = 1500, 950  # 窗口的宽度为 1000，高度为 600
        x = (screen_rect.width() - width) // 2  # 计算窗口水平方向的位置，使其居中。
        y = (screen_rect.height() - height) // 2 # 计算窗口垂直方向的位置，使其居中。
        self.setGeometry(x, y, width, height)  # 设置窗口的大小为 1000x600，并将其位置设置为 x, y。
        self.setFixedSize(width, height)  # 会将窗口的大小固定为 1000x600，这样用户就无法通过拖动窗口边界来改变窗口大小。

        # 创建一个垂直布局
        self.layout = QVBoxLayout()
        self.layout.setSpacing(15)  # 设置控件之间的间距为15像素

        # 创建一个QLabel来显示文字
        "在界面中间该界面文字信息代码"
        self.title_label = QLabel("基于YOLOv8目标检测系统", self)
        self.title_label.setAlignment(Qt.AlignCenter)  # 使文本居中显示
        self.layout.addWidget(self.title_label)  # 将文字标签添加到布局
        self.layout.addStretch(1)  # 在布局底部添加一个弹性空间，使文字标签保持在顶部
        # 设置文字大小和字体
        font = QFont()  # 创建一个字体对象
        font.setFamily("FangSong")  # 设置字体为 Arial（你可以根据需要换成其他字体）
        font.setPointSize(35)  # 设置字体大小为 36
        font.setBold(True)  # 设置加粗
        self.title_label.setFont(font)  # 将字体应用到 QLabel 上
        self.layout.addWidget(self.title_label)  # 将文字标签添加到布局

        # 将布局设置为窗口的布局
        self.setLayout(self.layout)

        # 创建并添加框架+
        self.result_frame = ResultFrame(self)
        self.operation_frame = OperationFrame(self)
        self.parameter_frame = ParameterFrame(self)
        self.backup_frame = BackupFrame(self)
        self.backup_frame = BackupFrame(self, self.run_detection)  # 将 run_detection 传递给 BackupFrame
        self.parameter_display = ParameterDisplay(self)

        self.is_running = True

        # 将信号连接到 ResultFrame 的 display_image_path_results 方法
        self.operation_frame.display_image.connect(self.result_frame.display_image_path_results)


    def run_detection(self):
        new_results = []
        def run(data_path, model_path='best.pt', conf_thres=0.5, iou_thres=0.45, imgsz=640, save_txt=False, save_conf=False, save_crop=False):
            """
            执行 YOLO 模型推理
            """
            # 加载训练好的模型
            model = YOLO(model_path)  # 使用指定路径的模型


            # 创建exp文件夹，并处理文件夹命名
            base_dir = 'runs'
            if not os.path.exists(base_dir):
                os.makedirs(base_dir)

            # 查找现有的exp文件夹
            exp_dir = os.path.join(base_dir, 'exp')
            if not os.path.exists(exp_dir):
                os.makedirs(exp_dir)
            else:
                # 如果exp文件夹已存在，生成类似exp1, exp2的文件夹
                i = 1
                while os.path.exists(os.path.join(base_dir, f'exp{i}')):
                    i += 1
                exp_dir = os.path.join(base_dir, f'exp{i}')  # 这里 f'exp{i}' 会被自动转换为字符串
                os.makedirs(exp_dir)


            if os.path.isdir(data_path):  # 文件夹路径
                image_files = [os.path.join(data_path, f) for f in os.listdir(data_path)
                               if f.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif',
                                                      '.tiff', '.webp', '.heif', '.raw', '.ico', '.jfif'))]
                video_files = [os.path.join(data_path, f) for f in os.listdir(data_path)
                               if f.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.flv', '.webm',
                                                      '.wmv', '.mpeg', '.mpg', '.3gp'))]

                # 如果需要保存txt文件，确保创建labels文件夹
                if save_txt:
                    labels_dir = os.path.join(exp_dir, 'labels')  # labels文件夹路径
                    if not os.path.exists(labels_dir):
                        os.makedirs(labels_dir)  # 创建labels文件夹

                # 如果需要保存裁剪图像，确保创建crops文件夹
                if save_crop:
                    crops_dir = os.path.join(exp_dir, 'crops')  # crops文件夹路径
                    if not os.path.exists(crops_dir):
                        os.makedirs(crops_dir)  # 创建crops文件夹

                # 处理图片文件
                for image_file in image_files:
                    # 获取原始图像
                    original_image = cv2.imread(image_file)
                    results = model.predict(image_file, conf=conf_thres, iou=iou_thres, imgsz=imgsz)[0]
                    im0 = results.plot()  # 获取带检测框的 numpy.ndarray 图像
                    self.result_frame.display_results(im0)  # 显示推理结果

                    # 保存检测结果图像
                    file_name = os.path.basename(image_file)  # 获取文件名
                    save_path = os.path.join(exp_dir, file_name)  # 组合保存路径
                    # 使用 OpenCV 保存图像
                    cv2.imwrite(save_path, im0)  # 这里 im0 是带有检测框的图像
                    new_results.append({"路径": save_path})

                    # 获取检测框的位置、置信度和类别
                    location_list = results.boxes.xyxy.tolist()  # [[x1, y1, x2, y2], ...]
                    conf_list = results.boxes.conf.tolist()  # [confidence1, confidence2, ...]
                    cls_list = results.boxes.cls.tolist()  # [class1, class2, ...]
                    name = results.names

                    # 如果需要保存txt文件
                    if save_txt:
                        # 获取对应的txt文件路径
                        txt_file_path = os.path.join(labels_dir, os.path.splitext(file_name)[0] + '.txt')

                        # 保存检测结果到txt文件
                        with open(txt_file_path, 'w') as f:
                            for location, conf, cls in zip(location_list, conf_list, cls_list):
                                # location = [x1, y1, x2, y2]
                                # 根据save_conf来决定是否保存置信度
                                if save_conf:
                                    f.write(f"{int(cls)} {' '.join(map(str, location))} {conf:.4f}\n")
                                else:
                                    f.write(f"{int(cls)} {' '.join(map(str, location))}\n")
                    if save_crop:
                        for i, (location, cls) in enumerate(zip(location_list, cls_list)):
                            # 获取类别名称，可以从模型中获取类名，假设模型有 classes 属性
                            class_name = name[int(cls)]  # 获取带检测框的 numpy.ndarray 图像

                            # 根据类别创建对应的文件夹
                            class_dir = os.path.join(crops_dir, class_name)
                            if not os.path.exists(class_dir):
                                os.makedirs(class_dir)

                            # 获取检测框的坐标
                            x1, y1, x2, y2 = map(int, location)

                            # 裁剪原始图像
                            cropped_image = original_image[y1:y2, x1:x2]  # [y1:y2, x1:x2]裁剪

                            # 保存裁剪后的图像
                            crop_file_name = f"{os.path.splitext(file_name)[0]}_{class_name}_{i}.jpg"
                            crop_save_path = os.path.join(class_dir, crop_file_name)

                            # 使用 OpenCV 保存裁剪的图像
                            cv2.imwrite(crop_save_path, cropped_image)


                # 处理视频文件
                def process_video(video_file):
                    cap = cv2.VideoCapture(video_file)
                    if not cap.isOpened():
                        print(f"无法打开视频文件 {video_file}")
                        return

                    # 获取视频的帧率和大小
                    fps = cap.get(cv2.CAP_PROP_FPS)  # 视频的帧率
                    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 视频的宽度
                    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 视频的高度

                    # 获取视频文件名并组合保存路径
                    file_name = os.path.basename(video_file)  # 获取原视频文件名
                    file_name_without_extension = os.path.splitext(file_name)[0]  # 去除文件扩展名
                    output_video_path = os.path.join(exp_dir, f"{file_name_without_extension}.mp4")  # 保存检测结果的视频路径

                    # 使用 VideoWriter 来保存处理后的视频
                    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 视频编码（使用 mp4 编码）
                    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))  # 输出视频对象

                    def display_video():
                        """视频推理并显示结果"""
                        frame_id = 0  # 帧计数器
                        while cap.isOpened() and self.is_running:
                            ret, frame = cap.read()  # 逐帧读取视频
                            if not ret:
                                break
                            # 推理单帧
                            result = model.predict(frame, conf=conf_thres, iou=iou_thres, imgsz=imgsz)[0]
                            im0 = result.plot()  # 获取带检测框的 numpy.ndarray 图像

                            # 获取检测框的位置、置信度和类别
                            location_list = result.boxes.xyxy.tolist()  # [[x1, y1, x2, y2], ...]
                            conf_list = result.boxes.conf.tolist()  # [confidence1, confidence2, ...]
                            cls_list = result.boxes.cls.tolist()  # [class1, class2, ...]
                            name = result.names

                            # 如果需要保存txt文件
                            if save_txt:
                                # 获取对应的txt文件路径，使用帧编号来命名
                                frame_txt_path = os.path.join(labels_dir,
                                                              f"{file_name_without_extension}_{frame_id}.txt")

                                # 保存检测结果到txt文件
                                with open(frame_txt_path, 'w') as f:
                                    for location, conf, cls in zip(location_list, conf_list, cls_list):
                                        # location = [x1, y1, x2, y2]
                                        # 根据save_conf来决定是否保存置信度
                                        if save_conf:
                                            f.write(f"{int(cls)} {' '.join(map(str, location))} {conf:.4f}\n")
                                        else:
                                            f.write(f"{int(cls)} {' '.join(map(str, location))}\n")

                            if save_crop:
                                for i, (location, cls) in enumerate(zip(location_list, cls_list)):
                                    # 获取类别名称，可以从模型中获取类名，假设模型有 classes 属性
                                    class_name = name[int(cls)]  # 获取带检测框的 numpy.ndarray 图像

                                    # 根据类别创建对应的文件夹
                                    class_dir = os.path.join(crops_dir, class_name)
                                    if not os.path.exists(class_dir):
                                        os.makedirs(class_dir)

                                    # 获取检测框的坐标
                                    x1, y1, x2, y2 = map(int, location)

                                    # 裁剪原始图像
                                    cropped_image = frame[y1:y2, x1:x2]  # [y1:y2, x1:x2]裁剪

                                    # 保存裁剪后的图像
                                    crop_file_name = f"{os.path.splitext(file_name)[0]}_{class_name}_{i}.jpg"
                                    crop_save_path = os.path.join(class_dir, crop_file_name)

                                    # 使用 OpenCV 保存裁剪的图像
                                    cv2.imwrite(crop_save_path, cropped_image)


                            # 将推理结果写入视频
                            out.write(im0)
                            # 显示推理结果
                            self.result_frame.display_results(im0)
                            cv2.waitKey(2)  # 等待 1ms，刷新显示
                            frame_id += 1  # 增加帧计数

                        # 释放资源
                        cap.release()
                        out.release()
                        cv2.destroyAllWindows()
                        print(f"保存的视频文件：{output_video_path}")
                        new_results.append({"路径": output_video_path})

                    # 启动视频推理线程
                    inference_thread = threading.Thread(target=display_video)
                    inference_thread.start()
                    inference_thread.join()  # 等待该线程完成再继续

                # 处理多个视频文件
                for video_file in video_files:
                    process_video(video_file)

            elif data_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.flv', '.webm', '.wmv', '.mpeg', '.mpg', '.3gp')):
                # 如果需要保存txt文件，确保创建labels文件夹
                if save_txt:
                    labels_dir = os.path.join(exp_dir, 'labels')  # labels文件夹路径
                    if not os.path.exists(labels_dir):
                        os.makedirs(labels_dir)  # 创建labels文件夹

                # 如果需要保存裁剪图像，确保创建crops文件夹
                if save_crop:
                    crops_dir = os.path.join(exp_dir, 'crops')  # crops文件夹路径
                    if not os.path.exists(crops_dir):
                        os.makedirs(crops_dir)  # 创建crops文件夹

                cap = cv2.VideoCapture(data_path)
                if not cap.isOpened():
                    print(f"无法打开视频文件 {data_path}")
                    return

                # 获取视频的帧率和大小
                fps = cap.get(cv2.CAP_PROP_FPS)  # 视频的帧率
                frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 视频的宽度
                frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 视频的高度

                # 获取视频文件名并组合保存路径
                file_name = os.path.basename(data_path)  # 获取原视频文件名
                file_name_without_extension = os.path.splitext(file_name)[0]  # 去除文件扩展名
                output_video_path = os.path.join(exp_dir,
                                                 f"{file_name_without_extension}.mp4")  # 保存检测结果的视频路径

                # 使用 VideoWriter 来保存处理后的视频
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 视频编码（使用 mp4 编码）
                out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))  # 输出视频对象

                def display_video():
                    """视频推理并显示结果"""
                    frame_id = 0  # 帧计数器
                    """视频推理并显示结果"""
                    while cap.isOpened() and self.is_running:
                        ret, frame = cap.read()  # 逐帧读取视频
                        if not ret:
                            break
                        # 推理单帧
                        result = model.predict(frame, conf=conf_thres, iou=iou_thres, imgsz=imgsz)[0]
                        im0 = result.plot()  # 获取带检测框的 numpy.ndarray 图像

                        # 获取检测框的位置、置信度和类别
                        location_list = result.boxes.xyxy.tolist()  # [[x1, y1, x2, y2], ...]
                        conf_list = result.boxes.conf.tolist()  # [confidence1, confidence2, ...]
                        cls_list = result.boxes.cls.tolist()  # [class1, class2, ...]
                        name = result.names


                        # 显示推理结果
                        self.result_frame.display_results(im0)  # 显示推理结果

                        # 如果需要保存txt文件
                        if save_txt:
                            # 获取对应的txt文件路径，使用帧编号来命名
                            frame_txt_path = os.path.join(labels_dir,
                                                          f"{file_name_without_extension}_{frame_id}.txt")

                            # 保存检测结果到txt文件
                            with open(frame_txt_path, 'w') as f:
                                for location, conf, cls in zip(location_list, conf_list, cls_list):
                                    # location = [x1, y1, x2, y2]
                                    # 根据save_conf来决定是否保存置信度
                                    if save_conf:
                                        f.write(f"{int(cls)} {' '.join(map(str, location))} {conf:.4f}\n")
                                    else:
                                        f.write(f"{int(cls)} {' '.join(map(str, location))}\n")

                        if save_crop:
                            for i, (location, cls) in enumerate(zip(location_list, cls_list)):
                                # 获取类别名称，可以从模型中获取类名，假设模型有 classes 属性
                                class_name = name[int(cls)]  # 获取带检测框的 numpy.ndarray 图像

                                # 根据类别创建对应的文件夹
                                class_dir = os.path.join(crops_dir, class_name)
                                if not os.path.exists(class_dir):
                                    os.makedirs(class_dir)

                                # 获取检测框的坐标
                                x1, y1, x2, y2 = map(int, location)

                                # 裁剪原始图像
                                cropped_image = frame[y1:y2, x1:x2]  # [y1:y2, x1:x2]裁剪

                                # 保存裁剪后的图像
                                crop_file_name = f"{os.path.splitext(file_name)[0]}_{class_name}_{i}.jpg"
                                crop_save_path = os.path.join(class_dir, crop_file_name)

                                # 使用 OpenCV 保存裁剪的图像
                                cv2.imwrite(crop_save_path, cropped_image)


                        # 将带检测框的图像写入到输出视频
                        out.write(im0)
                        cv2.waitKey(2)  # 等待 1ms，刷新显示
                        frame_id += 1  # 增加帧计数
                    cap.release()  # 释放视频捕获对象
                    out.release()  # 释放 VideoWriter 对象
                    cv2.destroyAllWindows()  # 关闭所有 OpenCV 窗口
                    new_results.append({"路径": output_video_path})
                # 启动一个新线程进行视频推理
                inference_thread = threading.Thread(target=display_video)
                inference_thread.start()


            elif data_path.isdigit():  # 摄像头 ID（数字形式）
                # 如果需要保存txt文件，确保创建labels文件夹
                if save_txt:
                    labels_dir = os.path.join(exp_dir, 'labels')  # labels文件夹路径
                    if not os.path.exists(labels_dir):
                        os.makedirs(labels_dir)  # 创建labels文件夹

                # 如果需要保存裁剪图像，确保创建crops文件夹
                if save_crop:
                    crops_dir = os.path.join(exp_dir, 'crops')  # crops文件夹路径
                    if not os.path.exists(crops_dir):
                        os.makedirs(crops_dir)  # 创建crops文件夹

                cap = cv2.VideoCapture(int(data_path))  # 通过摄像头 ID 打开摄像头
                if not cap.isOpened():
                    print(f"无法打开摄像头 {data_path}")
                    return

                # 获取视频的帧率和大小
                fps = cap.get(cv2.CAP_PROP_FPS)  # 视频的帧率
                frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 视频的宽度
                frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 视频的高度

                # 获取视频文件名并组合保存路径
                file_name = os.path.basename(data_path)  # 获取原视频文件名
                file_name_without_extension = os.path.splitext(file_name)[0]  # 去除文件扩展名
                output_video_path = os.path.join(exp_dir,
                                                 f"{file_name_without_extension}.mp4")  # 保存检测结果的视频路径

                # 使用 VideoWriter 来保存处理后的视频
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 视频编码（使用 mp4 编码）
                out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))  # 输出视频对象

                def display_video():
                    """视频推理并显示结果"""
                    frame_id = 0  # 帧计数器
                    """视频推理并显示结果"""
                    while cap.isOpened() and self.is_running:
                        ret, frame = cap.read()  # 逐帧读取视频
                        if not ret:
                            break
                        # 推理单帧
                        result = model.predict(frame, conf=conf_thres, iou=iou_thres, imgsz=imgsz)[0]
                        im0 = result.plot()  # 获取带检测框的 numpy.ndarray 图像

                        if not self.backup_frame.is_detecting:
                            print("检测停止")
                            self.backup_frame.is_detecting = True
                            return

                        # 获取检测框的位置、置信度和类别
                        location_list = result.boxes.xyxy.tolist()  # [[x1, y1, x2, y2], ...]
                        conf_list = result.boxes.conf.tolist()  # [confidence1, confidence2, ...]
                        cls_list = result.boxes.cls.tolist()  # [class1, class2, ...]
                        name = result.names

                        # 显示推理结果
                        self.result_frame.display_results(im0)  # 显示推理结果

                        # 如果需要保存txt文件
                        if save_txt:
                            # 获取对应的txt文件路径，使用帧编号来命名
                            frame_txt_path = os.path.join(labels_dir,
                                                          f"{file_name_without_extension}_{frame_id}.txt")

                            # 保存检测结果到txt文件
                            with open(frame_txt_path, 'w') as f:
                                for location, conf, cls in zip(location_list, conf_list, cls_list):
                                    # location = [x1, y1, x2, y2]
                                    # 根据save_conf来决定是否保存置信度
                                    if save_conf:
                                        f.write(f"{int(cls)} {' '.join(map(str, location))} {conf:.4f}\n")
                                    else:
                                        f.write(f"{int(cls)} {' '.join(map(str, location))}\n")

                        if save_crop:
                            for i, (location, cls) in enumerate(zip(location_list, cls_list)):
                                # 获取类别名称，可以从模型中获取类名，假设模型有 classes 属性
                                class_name = name[int(cls)]  # 获取带检测框的 numpy.ndarray 图像

                                # 根据类别创建对应的文件夹
                                class_dir = os.path.join(crops_dir, class_name)
                                if not os.path.exists(class_dir):
                                    os.makedirs(class_dir)

                                # 获取检测框的坐标
                                x1, y1, x2, y2 = map(int, location)

                                # 裁剪原始图像
                                cropped_image = frame[y1:y2, x1:x2]  # [y1:y2, x1:x2]裁剪

                                # 保存裁剪后的图像
                                crop_file_name = f"{os.path.splitext(file_name)[0]}_{class_name}_{i}.jpg"
                                crop_save_path = os.path.join(class_dir, crop_file_name)

                                # 使用 OpenCV 保存裁剪的图像
                                cv2.imwrite(crop_save_path, cropped_image)

                        # 将带检测框的图像写入到输出视频
                        out.write(im0)
                        cv2.waitKey(2)  # 等待 1ms，刷新显示
                        frame_id += 1  # 增加帧计数
                    cap.release()  # 释放视频捕获对象
                    out.release()  # 释放 VideoWriter 对象
                    cv2.destroyAllWindows()  # 关闭所有 OpenCV 窗口
                    new_results.append({"路径": output_video_path})

                # 启动一个新线程进行视频推理
                inference_thread = threading.Thread(target=display_video)
                inference_thread.start()


            else:  # 单张图片路径
                # 判断文件格式是否为图片
                if data_path.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp', '.jfif')):
                    try:
                        # 获取原始图像
                        original_image = cv2.imread(data_path)
                        results = model.predict(data_path, conf=conf_thres, iou=iou_thres, imgsz=imgsz)

                        # 如果需要保存txt文件，确保创建labels文件夹
                        if save_txt:
                            labels_dir = os.path.join(exp_dir, 'labels')  # labels文件夹路径
                            if not os.path.exists(labels_dir):
                                os.makedirs(labels_dir)  # 创建labels文件夹

                        # 如果需要保存裁剪图像，确保创建crops文件夹
                        if save_crop:
                            crops_dir = os.path.join(exp_dir, 'crops')  # crops文件夹路径
                            if not os.path.exists(crops_dir):
                                os.makedirs(crops_dir)  # 创建crops文件夹

                        for result in results:
                            im0 = result.plot()  # 获取带检测框的 numpy.ndarray 图像
                            self.result_frame.display_results(im0)  # 显示推理结果

                            location_list = result.boxes.xyxy.tolist()
                            conf_list = result.boxes.conf.tolist()
                            cls_list = result.boxes.cls.tolist()
                            name = result.names

                            # 保存检测结果图像
                            file_name = os.path.basename(data_path)  # 获取文件名
                            save_path = os.path.join(exp_dir, file_name)  # 组合保存路径
                            # 使用 OpenCV 保存图像
                            cv2.imwrite(save_path, im0)  # 这里 im0 是带有检测框的图像
                            new_results.append({"路径": save_path})

                            # 如果需要保存txt文件
                            if save_txt:
                                # 获取对应的txt文件路径
                                txt_file_path = os.path.join(labels_dir, os.path.splitext(file_name)[0] + '.txt')

                                with open(txt_file_path, 'w') as f:
                                    # 遍历每一个检测到的目标
                                    for location, conf, cls in zip(location_list, conf_list, cls_list):
                                        # location = [x1, y1, x2, y2]
                                        # 保存为 类别 置信度 x1 y1 x2 y2 格式，如果save_conf为True则包括置信度
                                        if save_conf:
                                            f.write(f"{int(cls)} {' '.join(map(str, location))} {conf:.4f}\n")
                                        else:
                                            f.write(f"{int(cls)} {' '.join(map(str, location))}\n")

                                # 如果需要裁剪目标
                            if save_crop:
                                for i, (location, cls) in enumerate(zip(location_list, cls_list)):
                                    # 获取类别名称，可以从模型中获取类名，假设模型有 classes 属性
                                    class_name = name[int(cls)]  # 获取带检测框的 numpy.ndarray 图像

                                    # 根据类别创建对应的文件夹
                                    class_dir = os.path.join(crops_dir, class_name)
                                    if not os.path.exists(class_dir):
                                        os.makedirs(class_dir)

                                    # 获取检测框的坐标
                                    x1, y1, x2, y2 = map(int, location)

                                    # 裁剪原始图像
                                    cropped_image = original_image[y1:y2, x1:x2]  # [y1:y2, x1:x2]裁剪

                                    # 保存裁剪后的图像
                                    crop_file_name = f"{os.path.splitext(file_name)[0]}_{class_name}_{i}.jpg"
                                    crop_save_path = os.path.join(class_dir, crop_file_name)

                                    # 使用 OpenCV 保存裁剪的图像
                                    cv2.imwrite(crop_save_path, cropped_image)

                    except Exception as e:
                        print(f"Error processing the image: {e}")
                else:
                    print(f"Provided path {data_path} is not a valid image file.")
            # 将多个检测结果传递给 `new_detection_result` 方法
            self.operation_frame.new_detection_result(new_results)



        def contains_chinese(path):
            # 判断路径中是否包含中文字符
            return bool(re.search(r'[\u4e00-\u9fff]', path))


        # 获取模型输入参数
        if self.parameter_frame.path_label4.text() == "开启":
            # 如果摄像头已开启
            data_path = "0"  # 这里假设"0"代表摄像头
        elif self.parameter_frame.path_label1.text():
            # 如果选择了图片
            data_path = self.parameter_frame.path_label1.text()
            if contains_chinese(data_path):
                msg = QMessageBox()
                msg.setIcon(QMessageBox.Warning)
                msg.setWindowTitle("输入路径有中文")
                msg.setText(f"图片路径有中文，请重新选择路径！\n当前路径：{data_path}")
                msg.setStandardButtons(QMessageBox.Ok)
                msg.exec_()
                return
        elif self.parameter_frame.path_label2.text():
            # 如果选择了文件夹
            data_path = self.parameter_frame.path_label2.text()
            if contains_chinese(data_path):
                msg = QMessageBox()
                msg.setIcon(QMessageBox.Warning)
                msg.setWindowTitle("输入路径有中文")
                msg.setText(f"文件夹路径有中文，请重新选择路径！\n当前路径：{data_path}")
                msg.setStandardButtons(QMessageBox.Ok)
                msg.exec_()
                return

            # 用一个列表记录包含中文字符的文件
            files_with_chinese = []
            # 遍历文件夹中的所有文件，检查文件名是否包含中文
            for filename in os.listdir(data_path):
                file_path = os.path.join(data_path, filename)
                # 只检查图片文件（根据文件扩展名）
                if os.path.isfile(file_path) and filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp', '.jfif')):
                    if contains_chinese(filename):
                        # 将包含中文的文件路径加入列表
                        files_with_chinese.append((filename, file_path))

            # 如果有中文文件，显示滚动的提示框
            if files_with_chinese:
                # 构建提示信息文本
                msg_text = "文件夹中包含中文的图片文件，请重新命名文件！\n\n"
                for filename, file_path in files_with_chinese:
                    # msg_text += f"当前文件：{filename}\n文件路径：{file_path}\n\n"
                    msg_text += f"{filename}\n"

                # 自定义对话框
                dialog = QDialog()
                dialog.setWindowTitle("图片文件名包含中文")
                layout = QVBoxLayout()

                # 创建 QTextEdit 来显示信息，并设置为只读
                text_edit = QTextEdit()
                text_edit.setText(msg_text)
                text_edit.setReadOnly(True)

                # 创建滚动区域，将 QTextEdit 放入其中
                scroll_area = QScrollArea()
                scroll_area.setWidget(text_edit)
                scroll_area.setWidgetResizable(True)

                # 创建确认按钮
                btn_ok = QPushButton("确认")
                btn_ok.clicked.connect(dialog.accept)  # 点击按钮关闭对话框

                # 添加控件到布局
                layout.addWidget(scroll_area)
                layout.addWidget(btn_ok)

                dialog.setLayout(layout)

                # 设置对话框的大小
                dialog.resize(600, 400)  # 调整对话框大小（宽600，高400）
                dialog.setFixedSize(dialog.size())  # 设置固定大小（防止调整）

                # 显示对话框
                dialog.exec_()
                return

        elif self.parameter_frame.path_label3.text():
            # 如果选择了视频
            data_path = self.parameter_frame.path_label3.text()
            if contains_chinese(data_path):
                msg = QMessageBox()
                msg.setIcon(QMessageBox.Warning)
                msg.setWindowTitle("输入路径有中文")
                msg.setText(f"视频路径有中文，请重新选择路径！\n当前路径：{data_path}")
                msg.setStandardButtons(QMessageBox.Ok)
                msg.exec_()
                return
        else:
            # 如果没有选择任何输入，弹出提示框
            msg = QMessageBox()
            msg.setIcon(QMessageBox.Warning)
            msg.setWindowTitle("输入数据提示")
            msg.setText("请输入模型的输入数据！\n数据路径不可以有中文")

            msg.setStandardButtons(QMessageBox.Ok)
            msg.exec_()  # 显示弹窗
            # 退出或其他操作
            return  # 或者你可以选择抛出异常或其他逻辑

        # # # 获取权重路径
        new_weights = self.parameter_display.path_label1.text()
        if self.parameter_display.path_label1.text() == "":
            # 如果没有选择任何输入，弹出提示框
            msg = QMessageBox()
            msg.setIcon(QMessageBox.Warning)
            msg.setWindowTitle("输入权重")
            msg.setText("请输入模型的权重！数据路径不可以有中文")
            msg.setStandardButtons(QMessageBox.Ok)
            msg.exec_()  # 显示弹窗
            # 退出或其他操作
            return  # 或者你可以选择抛出异常或其他逻辑

        if contains_chinese(new_weights):
            msg = QMessageBox()
            msg.setIcon(QMessageBox.Warning)
            msg.setWindowTitle("输入路径有中文")
            msg.setText(f"权重路径有中文，请重新选择路径！\n当前路径：{new_weights}")
            msg.setStandardButtons(QMessageBox.Ok)
            msg.exec_()
            return

        new_weights = Path(new_weights)
        model_path = new_weights

        # 获取置信度阈值
        new_confidence = self.parameter_display.path_label2.text()
        conf_thres = float(new_confidence)

        # 获取IOU阈值
        new_iou_thres = self.parameter_display.path_label3.text()
        iou_thres = float(new_iou_thres)

        # 获取图片大小参数
        new_imgsz = self.parameter_display.path_label4.text()
        imgsz = [int(new_imgsz), int(new_imgsz)]

        # # 是否保存txt结果
        if self.parameter_display.path_label5.text() == "是":
            # 保存txt结果
            save_txt = True
        else:
            save_txt = False

        # # 是否保存conf结果
        if self.parameter_display.path_label6.text() == "是":
            # 保存conf结果
            save_conf = True
        else:
            save_conf = False

        # 是否保存crop结果
        if self.parameter_display.path_label7.text() == "是":
            # 保存crop结果
            save_crop = True
        else:
            save_crop = False

        run(data_path=data_path,  model_path=model_path,
            conf_thres=conf_thres, iou_thres=iou_thres, imgsz=imgsz,
            save_txt=save_txt, save_conf=save_conf, save_crop=save_crop)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
